import random
from typing import Tuple


class RSA:
    def __init__(self, key_size: int = 1024):
        """
        初始化RSA实例
        :param key_size: 密钥长度，默认1024位
        """
        self.key_size = key_size
        self.public_key = None
        self.private_key = None

    def is_prime(self, n: int, k: int = 5) -> bool:
        """
        米勒-拉宾素性测试
        :param n: 要测试的数
        :param k: 测试次数
        :return: 是否为素数
        """
        if n < 2:
            return False
        if n in (2, 3):
            return True
        if n % 2 == 0:
            return False

        # 将n-1写成d*2^r的形式
        r, d = 0, n - 1
        while d % 2 == 0:
            r += 1
            d //= 2

        # 进行k次测试
        for _ in range(k):
            a = random.randint(2, n - 2)
            x = pow(a, d, n)
            if x == 1 or x == n - 1:
                continue

            for _ in range(r - 1):
                x = pow(x, 2, n)
                if x == n - 1:
                    break
            else:
                return False
        return True

    def generate_prime(self) -> int:
        """
        生成大素数
        :return: 素数
        """
        while True:
            # 生成一个奇数
            num = random.getrandbits(self.key_size // 2)
            num |= (1 << (self.key_size // 2 - 1)) | 1  # 确保是奇数且长度合适

            if self.is_prime(num):
                return num

    def gcd(self, a: int, b: int) -> int:
        """
        计算最大公约数
        """
        while b:
            a, b = b, a % b
        return a

    def mod_inverse(self, a: int, m: int) -> int:
        """
        计算模逆元
        :param a: 数
        :param m: 模数
        :return: a关于模m的逆元
        """

        def extended_gcd(a, b):
            if b == 0:
                return a, 1, 0
            gcd, x1, y1 = extended_gcd(b, a % b)
            x = y1
            y = x1 - (a // b) * y1
            return gcd, x, y

        gcd, x, _ = extended_gcd(a, m)
        if gcd != 1:
            raise ValueError("逆元不存在")
        return x % m

    def generate_keys(self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
        """
        生成RSA公钥和私钥
        :return: (公钥(n, e), 私钥(n, d))
        """
        # 生成两个大素数
        p = self.generate_prime()
        q = self.generate_prime()

        # 确保p和q不同
        while p == q:
            q = self.generate_prime()

        # 计算n和φ(n)
        n = p * q
        phi = (p - 1) * (q - 1)

        # 选择公钥指数e
        e = 65537  # 常用的公钥指数
        while self.gcd(e, phi) != 1:
            e = random.randint(2, phi - 1)

        # 计算私钥指数d
        d = self.mod_inverse(e, phi)

        self.public_key = (n, e)
        self.private_key = (n, d)

        return self.public_key, self.private_key

    def keys_to_string(self) -> Tuple[str, str]:
        """
        将公钥和私钥转换为字符串形式
        :return: (公钥字符串, 私钥字符串)
        """
        if not self.public_key or not self.private_key:
            raise ValueError("请先生成密钥对")

        public_key_str = f"{self.public_key[0]},{self.public_key[1]}"
        private_key_str = f"{self.private_key[0]},{self.private_key[1]}"

        return public_key_str, private_key_str

    def keys_from_string(self, public_key_str: str, private_key_str: str):
        """
        从字符串加载公钥和私钥
        :param public_key_str: 公钥字符串
        :param private_key_str: 私钥字符串
        """
        n_pub, e = map(int, public_key_str.split(','))
        n_priv, d = map(int, private_key_str.split(','))

        if n_pub != n_priv:
            raise ValueError("公钥和私钥不匹配")

        self.public_key = (n_pub, e)
        self.private_key = (n_priv, d)

    def encrypt(self, plaintext: str) -> str:
        """
        加密字符串
        :param plaintext: 明文
        :return: 密文字符串
        """
        if not self.public_key:
            raise ValueError("请先生成或加载公钥")

        n, e = self.public_key

        # 将字符串转换为字节，然后转换为整数
        plaintext_bytes = plaintext.encode('utf-8')
        plaintext_int = int.from_bytes(plaintext_bytes, 'big')

        # 确保明文小于n
        if plaintext_int >= n:
            raise ValueError("明文过长，请使用更长的密钥或分段加密")

        # 加密：c = m^e mod n
        ciphertext_int = pow(plaintext_int, e, n)

        # 将整数转换为十六进制字符串
        ciphertext_hex = hex(ciphertext_int)[2:]

        return ciphertext_hex

    def decrypt(self, ciphertext: str) -> str:
        """
        解密密文字符串
        :param ciphertext: 密文（十六进制字符串）
        :return: 明文字符串
        """
        if not self.private_key:
            raise ValueError("请先生成或加载私钥")

        n, d = self.private_key

        # 将十六进制字符串转换为整数
        ciphertext_int = int(ciphertext, 16)

        # 解密：m = c^d mod n
        plaintext_int = pow(ciphertext_int, d, n)

        # 将整数转换回字节
        # 计算所需的字节长度
        byte_length = (plaintext_int.bit_length() + 7) // 8
        plaintext_bytes = plaintext_int.to_bytes(byte_length, 'big')

        # 将字节解码为字符串
        plaintext = plaintext_bytes.decode('utf-8')

        return plaintext


# 使用示例
def main():
    # 创建RSA实例
    # 基本用法
    rsa = RSA(key_size=1024)
    # 生成密钥
    public_key, private_key = rsa.generate_keys()
    # 获取字符串形式的密钥
    pub_str, priv_str = rsa.keys_to_string()
    print("pub_str")
    print(pub_str)
    print("priv_str")
    print(priv_str)
    # 加密
    ciphertext = rsa.encrypt("Hello World!")
    print("ciphertext")
    print(ciphertext)
    # 解密
    plaintext = rsa.decrypt(ciphertext)
    print("plaintext")
    print(plaintext)


if __name__ == "__main__":
    main()
